import json
import os
import xml.etree.ElementTree as ET
from xml.dom import minidom

def create_voc_xml(json_file, output_dir):
    # 读取JSON文件
    with open(json_file, 'r', encoding='utf-8') as f:
        data = json.load(f)
    
    # 创建XML根元素
    root = ET.Element('annotation')
    
    # 添加基本信息
    folder = ET.SubElement(root, 'folder')
    folder.text = os.path.dirname(data['imagePath'])
    
    filename = ET.SubElement(root, 'filename')
    filename.text = data['imagePath']
    
    path = ET.SubElement(root, 'path')
    path.text = os.path.abspath(data['imagePath'])
    
    source = ET.SubElement(root, 'source')
    database = ET.SubElement(source, 'database')
    database.text = 'Unknown'
    
    # 添加图像尺寸信息
    size = ET.SubElement(root, 'size')
    width = ET.SubElement(size, 'width')
    width.text = str(data['imageWidth'])
    height = ET.SubElement(size, 'height')
    height.text = str(data['imageHeight'])
    depth = ET.SubElement(size, 'depth')
    depth.text = '3'  # 假设是RGB图像
    
    segmented = ET.SubElement(root, 'segmented')
    segmented.text = '0'
    
    # 添加每个目标框
    for shape in data['shapes']:
        obj = ET.SubElement(root, 'object')
        
        name = ET.SubElement(obj, 'name')
        name.text = shape['label']
        
        pose = ET.SubElement(obj, 'pose')
        pose.text = 'Unspecified'
        
        truncated = ET.SubElement(obj, 'truncated')
        truncated.text = '0'
        
        difficult = ET.SubElement(obj, 'difficult')
        difficult.text = '0'
        
        bndbox = ET.SubElement(obj, 'bndbox')
        xmin = ET.SubElement(bndbox, 'xmin')
        xmin.text = str(int(shape['points'][0][0]))
        ymin = ET.SubElement(bndbox, 'ymin')
        ymin.text = str(int(shape['points'][0][1]))
        xmax = ET.SubElement(bndbox, 'xmax')
        xmax.text = str(int(shape['points'][1][0]))
        ymax = ET.SubElement(bndbox, 'ymax')
        ymax.text = str(int(shape['points'][1][1]))
    
    # 将XML转换为字符串
    xml_str = minidom.parseString(ET.tostring(root)).toprettyxml(indent='    ')
    
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    # 生成输出文件名
    output_file = os.path.join(output_dir, os.path.splitext(data['imagePath'])[0] + '.xml')
    
    # 写入XML文件
    with open(output_file, 'w', encoding='utf-8') as f:
        f.write(xml_str)
    
    return output_file

def convert_directory(input_dir, output_dir):
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    # 遍历输入目录中的所有JSON文件
    for filename in os.listdir(input_dir):
        if filename.endswith('.json'):
            json_file = os.path.join(input_dir, filename)
            try:
                output_file = create_voc_xml(json_file, output_dir)
                print(f'已转换: {filename} -> {os.path.basename(output_file)}')
            except Exception as e:
                print(f'转换 {filename} 时出错: {str(e)}')

if __name__ == '__main__':
    # 设置输入和输出目录
    input_dir = '.'  # JSON文件所在目录
    output_dir = 'voc_annotations'  # 输出XML文件目录
    
    # 执行转换
    convert_directory(input_dir, output_dir) 